{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adfe5d3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#src pathing\n",
    "import os\n",
    "import sys\n",
    "\n",
    "import logging\n",
    "from os import path\n",
    "from typing import Optional, Tuple, Any, Callable, Dict, List, Optional\n",
    "from abc import ABC, abstractmethod\n",
    "from dataclasses import dataclass, field\n",
    "import random\n",
    "\n",
    "import argparse\n",
    "from types import SimpleNamespace\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "import torch.multiprocessing as mp\n",
    "import yfinance as yf # type: ignore\n",
    "\n",
    "\n",
    "from huggingface_hub import snapshot_download\n",
    "from torch.utils.data import Dataset\n",
    "from ffm import freq_map\n",
    "\n",
    "from tools.utils import print_model_statistics, make_logging_file, log_config\n",
    "\n",
    "from data_tools.TSdataset import TimeSeriesDataset_MultiCSV_train_Production, find_files_with_suffix\n",
    "from tools.model_utils import plot_predictions, get_model_FFM, FFM_weight_freeze\n",
    "from tools.inference_utils import plot_predictions_multi, plot_predictions_multi_distribution, plot_predictions_multi_distribution_v2\n",
    "\n",
    "\n",
    "from ffm.ffm_base import FFmHparams"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86868081",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dataclasses\n",
    "from dataclasses import dataclass, field\n",
    "\n",
    "@dataclasses.dataclass(kw_only=True)\n",
    "class Freq_map_dict:\n",
    "    major_6_bench_map: dict = field(\n",
    "        default_factory=lambda:\n",
    "    {\n",
    "    \"ettm2\": \"15min\",\n",
    "    \"ettm1\": \"15min\",\n",
    "    \"etth2\": \"H\",\n",
    "    \"etth1\": \"H\",\n",
    "    \"electricity\": \"H\",\n",
    "    \"traffic\": \"H\",\n",
    "    \"weather\": \"10min\",\n",
    "    \"national_illness\": 'W',\n",
    "    \"exchange_rate\": 'D',\n",
    "    }\n",
    "    )\n",
    "    \n",
    "    \n",
    "    major_6_bench_val_map: dict = field(\n",
    "        default_factory=lambda:\n",
    "    {\n",
    "    \"val_elec\": \"H\",\n",
    "    \"val_etth1\": \"H\",\n",
    "    \"val_ettm1\": \"15min\",\n",
    "    \"val_exchange\": \"D\",\n",
    "    \"val_illness\": \"W\",\n",
    "    \"val_traffic\": \"H\",\n",
    "    \"val_weather\": \"10min\",\n",
    "    }\n",
    "    )\n",
    "\n",
    "    universal_map: dict = field(\n",
    "        default_factory=lambda:\n",
    "    {\n",
    "    \"_1d.csv\": 0,\n",
    "    \"_1wk.csv\": 1,\n",
    "    \"_1h.csv\": 0,\n",
    "    \"_1m.csv\": 0,\n",
    "    }\n",
    "    )\n",
    "\n",
    "\n",
    "FREQ_POSSIBLE_CONTEXT_LENGTH = {\n",
    "    0 : [512, 256, 128],\n",
    "    1 : [256, 128],\n",
    "    2 : [64],\n",
    "}\n",
    "\n",
    "\n",
    "DATA_SLICE_INTERVAL_SMALL_D = {\n",
    "    0 : 32,\n",
    "    1 : 1,\n",
    "    2 : 1,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d6bc413",
   "metadata": {},
   "outputs": [],
   "source": [
    "config = SimpleNamespace()\n",
    "\n",
    "# plt setting (if any)\n",
    "config.random_seed = random.randint(0, 100000)\n",
    "config.save_dir = r'pics/vis2'\n",
    "config.save_name = ''\n",
    "\n",
    "# data datasets\\val_datasets_ts_major6  datasets\\stock_v1\\val_v1_nv datasets\\stock_v1\\val_v1_nv_m_h  datasets\\stock_v1\\test_v1_nv_flat\n",
    "config.data_folder = r'datasets\\stock_v1\\test_v1_nv_flat'  # replace with your actual path\n",
    "config.num_workers = 2\n",
    "config.series_norm = False\n",
    "config.mask_ratio = 0\n",
    "config.freq_map_mode = 1             #0 is direct conversion, 1 is custom suffix name match\n",
    "\n",
    "\n",
    "\n",
    "# training hyperparams\n",
    "config.batch_size = 64\n",
    "\n",
    "# model change param on this\n",
    "config.checkpoint = r'checkpoints\\hbcloss_ft_v4_lowquantile\\hbc_v4_ep3_lowq_trained.pth'  # replace with actual path\n",
    "config.num_experts = 4\n",
    "config.gating_top_n = 2\n",
    "config.load_from_compile = True\n",
    "\n",
    "# device\n",
    "config.device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "config.gpu_ids = [0]\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d86f89b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "#model loading\n",
    "ffm_hparams = FFmHparams(num_experts=config.num_experts,\n",
    "                        gating_top_n=config.gating_top_n,\n",
    "                        load_from_compile=config.load_from_compile,)\n",
    "\n",
    "model, ffm_config, ffm_api = get_model_FFM(checkpoint=config.checkpoint, hparams=ffm_hparams)\n",
    "\n",
    "print_model_statistics(model=model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22b2c87d",
   "metadata": {},
   "outputs": [],
   "source": [
    " #datasets\n",
    "config.context_length_list = FREQ_POSSIBLE_CONTEXT_LENGTH #[32, 64, 128 ,256 , 512, 1024, 2048] #context length for variable length input\n",
    "config.data_slice_interval = DATA_SLICE_INTERVAL_SMALL_D\n",
    "freq_map_dict = Freq_map_dict()\n",
    "train_freq_map = {}\n",
    "eval_freq_map = {}\n",
    "\n",
    "if config.freq_map_mode == 0:\n",
    "    for k, v in freq_map_dict.major_6_bench_map.items():\n",
    "        train_freq_map[k] = freq_map(v)\n",
    "    \n",
    "    for k, v in freq_map_dict.major_6_bench_val_map.items():\n",
    "        eval_freq_map[k] = freq_map(v)\n",
    "elif config.freq_map_mode == 1:\n",
    "    train_freq_map = freq_map_dict.universal_map\n",
    "    eval_freq_map = freq_map_dict.universal_map\n",
    "\n",
    "\n",
    "#set up plt dataset\n",
    "val_file_list = find_files_with_suffix(directory=config.data_folder, suffix='.csv')\n",
    "val_dataset = TimeSeriesDataset_MultiCSV_train_Production(csv_paths=val_file_list, horizon_length=FFmHparams.output_patch_len,\n",
    "                                                freq_map=eval_freq_map, freq_map_mode=config.freq_map_mode,\n",
    "                                                mask_ratio=config.mask_ratio,\n",
    "                                                possible_context_lengths=config.context_length_list,\n",
    "                                                series_norm=config.series_norm,\n",
    "                                                data_slice_interval=config.data_slice_interval,\n",
    "                                                shuffle_seed=config.random_seed\n",
    "                                                )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23e87c67",
   "metadata": {},
   "outputs": [],
   "source": [
    "#output\n",
    "output_length = 60\n",
    "quantile_list = [1, 3, 7, 9]              #0 for mean, 1 - 9 for quantiles with 1 increment\n",
    "config.num_img = 50\n",
    "output_number=True\n",
    "\n",
    "plot_predictions_multi_distribution_v2(\n",
    "    model=model,\n",
    "    val_dataset=val_dataset,\n",
    "    number_img=config.num_img,\n",
    "    model_version=os.path.basename(config.checkpoint),\n",
    "    save_dir=None, #config.save_dir,\n",
    "    save_name=config.save_name,\n",
    "    moe_n=config.num_experts,\n",
    "    moe_tk=config.gating_top_n,\n",
    "    quantiles=quantile_list,\n",
    "    output_length=output_length,\n",
    "    output_number=output_number,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d21f7223",
   "metadata": {},
   "outputs": [],
   "source": [
    "#output\n",
    "output_length = 96\n",
    "quantile = 0              #0 for mean, 1 - 9 for quantiles with 1 increment\n",
    "config.num_img = 20\n",
    "\n",
    "plot_predictions_multi(\n",
    "    model=model,\n",
    "    val_dataset=val_dataset,\n",
    "    number_img=config.num_img,\n",
    "    model_version=os.path.basename(config.checkpoint),\n",
    "    save_dir=None, #config.save_dir,\n",
    "    save_name=config.save_name,\n",
    "    moe_n=config.num_experts,\n",
    "    moe_tk=config.gating_top_n,\n",
    "    quantile=quantile,\n",
    "    output_length=output_length,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d38350d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "#output\n",
    "output_length = 96\n",
    "quantile = 5             #0 for mean, 1 - 9 for quantiles with 1 increment\n",
    "config.num_img = 20\n",
    "\n",
    "plot_predictions_multi(\n",
    "    model=model,\n",
    "    val_dataset=val_dataset,\n",
    "    number_img=config.num_img,\n",
    "    model_version=os.path.basename(config.checkpoint),\n",
    "    save_dir=None, #config.save_dir,\n",
    "    save_name=config.save_name,\n",
    "    moe_n=config.num_experts,\n",
    "    moe_tk=config.gating_top_n,\n",
    "    quantile=quantile,\n",
    "    output_length=output_length,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b334ddf",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gc\n",
    "del model  # Delete the model reference\n",
    "gc.collect()  # Collect garbage\n",
    "torch.cuda.empty_cache()  # Clear cached memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d6914b6",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
